from matplotlib.pyplot import figure, show
import matplotlib.animation as animation
from numpy import zeros, arange, linspace, min, max, mean
from h5py import File
from afd.plots.cylinder import bit_map
from pdb import set_trace


def composite_animation(r, z, t, T, vL):
    """
    Create a composite figure of the simulation data.
    r : [nr x nz] numpy array of the radial coordinates
    z : [nr x nz] numpy array of the axial coordinates
    T : [nt x nr x nz] numpy array of temperature data
    vL : [nt x nr x nz] numpy array of liquid volume fraction data
    return FuncAnimation   Note: Must return the FuncAnimation function or only
        1 frame will be drawn before the function is discarded with garbage
        collection!
    """
    def _initialize():
        """
        Function to initialize plots called by matplotlib
        animation.FuncAnimimation.
        return list or tuple of matplotlib objects. (see Note on _update regarding return)
        """
        plots[0].set_data(zeros((320,400,3)))
        plots[1].set_data(zeros((320,400,3)))
        #for c in plots[2].collections:
        #    ax[2].collections.remove(c)
        #plots[2] = ax[2].contourf(r, z, T[0], linspace(min(T[0])-1, max(T[0])+1))
        #for c in plots[3].collections:
        #    ax[3].collections.remove(c)
        #plots[3] = ax[3].contourf(r, z, vL[0], linspace(min(vL[0])-.01, max(vL[0])+.01))
        plots[2].set_data([],[])
        plots[3].set_data([],[])
        return plots

    def _update(i):
        """
        Function used to update the data for the individual plots.
        Called by the matplotlib animation.FuncAnimation function.
        i : integer number of times the update function has been called
            Note: the first argument of update (i) is automatically supplied by FuncAnimation,
            all other arguments are passed according to the keyword, fargs, arguments
            defined by the FuncAnimation call.
        return list or tuple of matplotlib objects
        """
        plots[0].set_data(T_cyl.next())
        plots[1].set_data(vL_cyl.next())
        #plots[4].set_data(arange(0,i),mean(mean(T[0:i],axis=2),axis=1))
        #for c in plots[2].collections:
        #    ax[2].collections.remove(c)
        #plots[2] = ax[2].contourf(r, z, T[i], linspace(min(T[i])-1, max(T[i])+1))
        #for c in plots[3].collections:
        #    ax[3].collections.remove(c)
        #plots[3] = ax[3].contourf(r, z, vL[i], linspace(min(vL[i])-.01, max(vL[i])+.01))
        plots[2].set_data(t[:i],mean(mean(T[0:i],axis=2),axis=1))
        plots[3].set_data(t[:i],mean(mean(vL[0:i],axis=2),axis=1))
        return plots

    def setup(T_cyl, vL_cyl):
        fig = figure()
        ax = [fig.add_subplot(2,2,i) for i in range(1,5)]
        plots = [ax[0].imshow(T_cyl.next())]
        ax[0].set_title('temperature')
        plots.append(ax[1].imshow(vL_cyl.next()))
        ax[1].set_title('liquid')
        #plots.append(ax[2].contourf(r, z, T[0], linspace(min(T[0])-1, max(T[0])+1)))
        #plots.append(ax[3].contourf(r, z, vL[0], linspace(min(vL[0])-.01, max(vL[0])+0.01)))
        #plots.append(ax[4].plot([],[]))
        plots.append(ax[2].plot([],[])[0])
        ax[2].set_xlim((t.min(), t.max()))
        ax[2].set_ylim((T.min(), T.max()))
        ax[2].set_xlabel(r'time (s)')
        ax[2].set_ylabel(r'temperature (K)')
        plots.append(ax[3].plot([],[])[0])
        ax[3].set_xlim((t.min(), t.max()))
        ax[3].set_ylim((vL.min(), vL.max()))
        ax[3].set_xlabel(r'time (s)')
        ax[3].set_ylabel(r'$\phi_L$')
        return fig, ax, plots

    
    T_cyl = bit_map(r, z, 25, T, "Temperature")
    vL_cyl = bit_map(r, z, 25, vL, "Liquid")

    fig, ax, plots = setup(T_cyl, vL_cyl)

    """
    fig : figure to update
    update : function to call that updates the data
    interval : time in ms between calls to update
    blit : boolean to redraw all (blit=False) or just changed (blit=True) pixels
    """
    anim = animation.FuncAnimation(fig, _update, init_func=_initialize, frames=t.size-1, interval=10, blit=True)
    show()
    anim.save('test.mp4')
    return anim

def test():
    """
    Retrieve data from hdf5 file for test purposes
    """
    f_h5py = File('simulation results.hdf5', mode='r')
    t = f_h5py['time step/1e-4/t'].value
    r = f_h5py['time step/1e-4/r'].value
    z = f_h5py['time step/1e-4/z'].value
    T = f_h5py['time step/1e-4/T'].value
    vL = f_h5py['time step/1e-4/v_L'].value
    f_h5py.close()
    try:
        anim = composite_animation(r, z, t[0:3000:10], T[0:3000:10], vL[0:3000:10])
    except:
        pass
    return anim

